"""
Hard Parameter Sharing Architecture
硬参数共享架构实现
"""

from typing import Dict, List, Any, Optional
from ml_lib.core import Tensor
from ml_lib.nn.module import Module
from .base import MTLArchitecture


class HardParameterSharing(MTLArchitecture):
    """
    硬参数共享架构
    
    用户需要提供自定义的共享编码器和任务特定头部，
    框架负责组织前向传播和参数管理。
    """
    
    def __init__(self, input_dim: int, task_configs: Dict[str, Any], 
                 shared_encoder: Module, task_heads: Dict[str, Module]):
        """
        初始化硬参数共享架构
        
        Args:
            input_dim: 输入维度
            task_configs: 任务配置字典
            shared_encoder: 用户自定义的共享编码器
            task_heads: 用户自定义的任务特定头部字典 {task_name: head_module}
        """
        super().__init__(input_dim, task_configs)
        
        # 验证任务头部是否完整
        for task_name in self.task_names:
            if task_name not in task_heads:
                raise ValueError(f"缺少任务 {task_name} 的头部定义")
        
        # 存储用户提供的组件
        self.shared_encoder = shared_encoder
        self.task_heads = task_heads
        
        # 将任务头部注册为模块参数
        for task_name, head in task_heads.items():
            setattr(self, f'head_{task_name}', head)
    
    def forward(self, x: Tensor) -> Dict[str, Tensor]:
        """
        前向传播
        
        Args:
            x: 输入张量 [batch_size, input_dim]
            
        Returns:
            每个任务的输出字典
        """
        # 共享特征提取
        shared_features = self.shared_encoder(x)
        
        # 任务特定预测
        outputs = {}
        for task_name in self.task_names:
            outputs[task_name] = self.task_heads[task_name](shared_features)
        
        return outputs
    
    def get_shared_representation(self, x: Tensor) -> Tensor:
        """
        获取共享表示
        
        Args:
            x: 输入张量
            
        Returns:
            共享表示张量
        """
        return self.shared_encoder(x)
    
    def get_task_specific_params(self, task_name: str) -> List[Any]:
        """获取特定任务的参数"""
        if task_name not in self.task_heads:
            raise ValueError(f"任务 {task_name} 不存在")
        
        return list(self.task_heads[task_name].parameters())
    
    def get_shared_params(self) -> List[Any]:
        """获取共享参数"""
        return list(self.shared_encoder.parameters())
    
    def __repr__(self):
        return (f"HardParameterSharing(\n"
                f"  input_dim={self.input_dim},\n"
                f"  shared_encoder={type(self.shared_encoder).__name__},\n"
                f"  num_tasks={self.num_tasks},\n"
                f"  tasks={self.task_names}\n"
                f")") 